import time
import os.path
import tensorflow as tf
from GA import GA_Util
from ResNet import ResNet
from data_providers.utils import get_data_provider_by_name


N_CLASSES = 10
BATCH_SIZE = 50
MOMENTUM = 0.9
WEIGHT_DECAY = 0.0001
WARMUP_STEP = 20
BUILD_BLOCKS = 9
GEN_SIZE = BUILD_BLOCKS * 3
SP_SET = [0, BUILD_BLOCKS, BUILD_BLOCKS*2]
MODEL_NAME = 'ResNet-56'

train_params = {
    'batch_size': BATCH_SIZE,
    'validation_set': True,
    'validation_split': 0.1,
    'shuffle': 'every_epoch',
    'normalization': 'by_chanels',
    'save_path': os.path.join('..', 'data', 'cifar')
}

data_provider = get_data_provider_by_name('C10+', train_params)

def wp_eval_func(ori_acc, target='B'):
    def eval_func(state_switch):
        hp = {
            'learning_rate': 0.1,
            'keep_prob': 1.0,
            'weight_decay': WEIGHT_DECAY,
            'momentum': MOMENTUM,
        }
        if target == 'B':   #如果target==B，是对block进行剪枝，否则是对filter进行剪枝。这个代码是粗粒度剪枝细粒度剪枝合一起的
            hp['block_ss'] = state_switch
            hp['filter_ss'] = [1]*model.ss_size
        else:
            hp['block_ss'] = [1]*GEN_SIZE
            hp['filter_ss'] = state_switch
        result = model.val_test(data_provider, hp)
        return result[0]*100.0/ori_acc
    return eval_func

def warmup_training(global_epoch, total_epoch):#返回测试正确率
    filter_size = model.ss_size#filter的总数
    warmup_hyper_param = {
        'learning_rate': 0.1,
        'keep_prob': 1.0,
        'weight_decay': WEIGHT_DECAY,
        'momentum': MOMENTUM,
        'block_ss': [1]*GEN_SIZE,
        'filter_ss': [1]*filter_size
    }
    for x in filter_drop:
        warmup_hyper_param['filter_ss'][x] = 0
    for i in range(1, total_epoch+1):
        st_time = time.time()
        acc, los = model.train_single_epoch(data_provider, warmup_hyper_param)
        print('%dth train-epoch ACC:%.2f LOS:%.6f' % (i+global_epoch, acc*100.0, los))
        en_time = time.time()
        print('*%dth epoch train_time:%d*' % (i+global_epoch, int((en_time-st_time)*1000)))
        if i == int(total_epoch*0.8):
            warmup_hyper_param['learning_rate'] = 0.05
    acc, los = model.val_test(data_provider, warmup_hyper_param)
    print('%dth valia-epoch ACC:%.2f LOS:%.6f' % (global_epoch+total_epoch, acc*100.0, los))
    return acc*100.0

def normal_training(global_epoch):
    filter_size = model.ss_size
    normal_hp = {
        'learning_rate': 0.1,
        'keep_prob': 1.0,
        'weight_decay': WEIGHT_DECAY,
        'momentum': MOMENTUM,
        'block_ss': [1]*GEN_SIZE,
        'filter_ss': [1]*filter_size
    }
    for x in filter_drop:
        normal_hp['filter_ss'][x] = 0
    cnt = 0
    while (cnt < 10) or (global_epoch < 120):
        cnt += 1
        global_epoch += 1
        st_time = time.time()
        acc, los = model.train_single_epoch(data_provider, normal_hp)
        print('%dth train-epoch ACC:%.2f LOS:%.6f' % (global_epoch, acc*100.0, los))
        en_time = time.time()
        print('*%dth epoch train_time:%d*' % (global_epoch, int((en_time-st_time)*1000)))
    acc, los = model.val_test(data_provider, normal_hp)
    print('%dth valia-epoch ACC:%.2f LOS:%.6f' % (global_epoch, acc*100.0, los))
    normal_hp['learning_rate'] = 0.01
    for i in range(1, 121):
        st_time = time.time()
        acc, los = model.train_single_epoch(data_provider, normal_hp)
        print('%dth train-epoch ACC:%.2f LOS:%.6f' % (global_epoch+i, acc*100.0, los))
        en_time = time.time()
        print('*%dth epoch train_time:%d*' % (global_epoch+i, int((en_time-st_time)*1000)))
        if i == 40:
            normal_hp['learning_rate'] = 0.001
            acc, los = model.val_test(data_provider, normal_hp)
            print('%dth valia-epoch ACC:%.2f LOS:%.6f' % (global_epoch+40, acc*100.0, los))
        if i == 80:
            normal_hp['learning_rate'] = 0.0002
            acc, los = model.val_test(data_provider, normal_hp)
            print('%dth valia-epoch ACC:%.2f LOS:%.6f' % (global_epoch+80, acc*100.0, los))
    acc, los = model.val_test(data_provider, normal_hp)
    print('%dth valia-epoch ACC:%.2f LOS:%.6f' % (global_epoch+120, acc*100.0, los))

block_drop = []
filter_drop = []
model = ResNet(MODEL_NAME, BUILD_BLOCKS, N_CLASSES, BATCH_SIZE, drop_set=block_drop)

global_epoch = 0
acc = warmup_training(global_epoch, WARMUP_STEP)
global_epoch += WARMUP_STEP
model.saver.save(model.sess, './model/resnet56-ori')
st_time = time.time()
ea_helper = GA_Util(
    'es_ea',
    30,
    GEN_SIZE,
    eval_fun=wp_eval_func(acc, target='B'),
    max_gen=25,
    drop_set=block_drop,
    beta=1.8,
    sp_set=SP_SET
)
elite = ea_helper.evolution()
elite = ea_helper.population[elite]
en_time = time.time()
print(elite)
for x in range(GEN_SIZE):
    if elite[x] == 0:
        block_drop.append(x)
print('Drop Set:', block_drop)
print('Drop Count:', len(block_drop))
print('GA Consume:', int((en_time-st_time)*1000))

tf.reset_default_graph()
model = ResNet(MODEL_NAME, BUILD_BLOCKS, N_CLASSES, BATCH_SIZE, drop_set=block_drop)
model.saver.restore(model.sess, './model/resnet56-ori')


acc = warmup_training(global_epoch, WARMUP_STEP)
global_epoch += WARMUP_STEP
model.saver.save(model.sess, './model/resnet-finetune-after-blockdrop')
st_time = time.time()
ea_helper = GA_Util(
    'es_ea',
    30,
    model.ss_size,
    eval_fun=wp_eval_func(acc, target='F'),
    max_gen=50,
    drop_set=filter_drop,
    beta=2.0,
    sp_set=[],
    target='F'
)
elite = ea_helper.evolution()
elite = ea_helper.population[elite]
en_time = time.time()
print(elite)
filter_drop.clear()
for x in range(model.ss_size):
    if elite[x] == 0:
        filter_drop.append(x)
print('Drop Set:', filter_drop)
print('Drop Count:', len(filter_drop))
print('GA Consume:', int((en_time-st_time)*1000))

normal_training(global_epoch)
model.saver.save(model.sess, './model/resnet-final-pruned')

b_ss = [1] * GEN_SIZE
f_ss = [1] * model.ss_size

for x in filter_drop:
    f_ss[x] = 0

hyper_param = {
    'weight_decay': WEIGHT_DECAY,
    'block_ss': b_ss,
    'filter_ss': f_ss
}

acc, los = model.model_test(data_provider, hyper_param)
print('Final Acc:%.2f; Final Los:%.6f'%(acc*100.0, los))